import os
import json
from typing import List, Dict
import re
from openai import OpenAI

class LongMemoryAssistant:
    def __init__(self, api_key: str, model: str = "deepseek-reasoner"):
        self.api_key = api_key
        self.model = model
        self.max_history = 20  # 保留最近20轮对话
        self.dialogue_history: List[Dict] = []
        self.sensitive_words = ["自杀", "诈骗", "政治"]  # 敏感词过滤列表
        self.client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com/v1")

    def add_message(self, role: str, content: str):
        """添加对话记录"""
        self.dialogue_history.append({"role": role, "content": content})
        # 控制历史长度
        if len(self.dialogue_history) > self.max_history * 2:
            self.dialogue_history = self.dialogue_history[-self.max_history * 2:]

    def generate_response(self, user_input: str) -> str:
        """生成回复（核心逻辑）"""
        # 1. 添加用户输入
        self.add_message("user", user_input)

        # 2. 调用大模型API
        try:
            messages = self._build_context()
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages
            )
            reply = response.choices[0].message.content
        except Exception as e:
            reply = f"请求失败：{str(e)}"

        # 3. 安全过滤
        if self._safety_check(reply):
            reply = "此问题需要人工介入处理，请稍候。"

        # 4. 添加助理回复
        self.add_message("assistant", reply)
        return reply

    def _build_context(self) -> List[Dict]:
        """构建带长上下文记忆的prompt"""
        # 示例：自动保留关键实体（如订单号、法律条款）
        context = []
        for msg in self.dialogue_history[-10:]:  # 取最近10轮
            context.append({"role": msg["role"], "content": msg["content"]})

        # 添加系统提示词（法律/客服领域增强）
        system_prompt = """
        你是一名专业律师，回答需满足：
        1. 引用《劳动合同法》等法律条款
        2. 分步骤给出建议
        3. 用通俗语言解释专业术语
        """
        context.insert(0, {"role": "system", "content": system_prompt})
        return context

    def _safety_check(self, text: str) -> bool:
        """安全审查"""
        # 敏感词检测
        if any(word in text for word in self.sensitive_words):
            return True
        # 合规性正则检测（示例：防止生成联系方式）
        if re.search(r"\d{11}|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}", text):
            return True
        return False

    def save_conversation(self, file_path: str):
        """保存对话记录"""
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(self.dialogue_history, f, ensure_ascii=False, indent=2)


# ----------------- 使用示例 -----------------
if __name__ == "__main__":
    # 初始化（需替换实际API KEY）
    api_key = "sk-241d627b84654fa2b11610a7498d0162"
    if api_key is None:
        print("请设置有效的API KEY环境变量。")
    else:
        assistant = LongMemoryAssistant(api_key=api_key)

        # 模拟多轮对话
        conversations = [
            "公司突然解雇我且没有书面通知，合法吗？",
            "我入职3年，月薪2万元",
            "仲裁需要准备哪些材料？"
        ]

        for query in conversations:
            print(f"用户：{query}")
            response = assistant.generate_response(query)
            print(f"助理：{response}\n")

        # 保存对话记录
        assistant.save_conversation("legal_conversation.json")